from __future__ import annotations

import argparse
import os
from typing import List

import torch
from PIL import Image
from torchvision import transforms

from mnist_model import load_model


def load_images(image_path: str) -> List[str]:
    if os.path.isdir(image_path):
        files = [os.path.join(image_path, f) for f in os.listdir(image_path) if f.lower().endswith((".png", ".jpg", ".jpeg", ".bmp"))]
        files.sort()
        return files
    else:
        return [image_path]


def build_transform() -> transforms.Compose:
    return transforms.Compose([
        transforms.Grayscale(num_output_channels=1),
        transforms.Resize((28, 28)),
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])


def predict(model_path: str, image_path: str, device: str = "cpu") -> None:
    device_t = torch.device(device)
    model = load_model(model_path, device=device_t)
    model.eval()

    preprocessing = build_transform()
    image_files = load_images(image_path)

    for fpath in image_files:
        image = Image.open(fpath).convert("L")
        tensor = preprocessing(image).unsqueeze(0).to(device_t)
        with torch.no_grad():
            logits = model(tensor)
            probs = torch.softmax(logits, dim=1)
            pred = int(probs.argmax(dim=1).item())
            confidence = float(probs.max().item())
        print(f"{os.path.basename(fpath)} => {pred} (conf {confidence:.4f})")


def main() -> None:
    parser = argparse.ArgumentParser(description="Predict digits with trained MNIST model")
    parser.add_argument("--weights", type=str, required=True, help="Path to model weights (e.g., outputs/mnist_cnn_best.pt)")
    parser.add_argument("--image", type=str, required=True, help="Image file or directory")
    parser.add_argument("--device", type=str, default="cpu")
    args = parser.parse_args()

    predict(args.weights, args.image, device=args.device)


if __name__ == "__main__":
    main()